# c.py
import matplotlib.pyplot as plt
import networkx as nx
import matplotlib.animation as animation
import random


def graph_coloring_visualization(ax):
    # 随机生成节点数，范围在[5, 10]之间
    num_nodes = random.randint(5, 10)

    # 创建一个随机图，保证连通性
    G = nx.gnm_random_graph(num_nodes, random.randint(num_nodes, num_nodes * (num_nodes - 1) // 2))
    pos = nx.spring_layout(G)

    # 使用贪心算法进行图着色并返回使用的颜色数
    def greedy_coloring(graph):
        color_map = {}
        colors = ['red', 'green', 'blue', 'yellow']  # 最多使用四种颜色
        for node in graph.nodes():
            available_colors = set(colors)
            for neighbor in graph.neighbors(node):
                if neighbor in color_map:
                    available_colors.discard(color_map[neighbor])
            if available_colors:
                color_map[node] = available_colors.pop()
            else:
                color_map[node] = colors[0]  # fallback to first color, should not happen in a planar graph
        used_colors = set(color_map.values())
        return color_map, len(used_colors)

    color_map, used_colors = greedy_coloring(G)

    # 初始化所有节点为灰色
    node_colors = ['gray'] * len(G.nodes())

    def update(num):
        ax.clear()
        # 更新节点颜色
        if num < len(node_colors):
            node_colors[list(G.nodes())[num]] = color_map[list(G.nodes())[num]]
        nx.draw(G, pos, with_labels=True, node_color=node_colors, ax=ax)
        ax.set_title(f"Step {num + 1}")

    ani = animation.FuncAnimation(ax.figure, update, frames=len(G.nodes()), repeat=False, interval=1000)

    # 在图表外部显示节点数和使用的最少颜色数
    plt.figtext(0.1, 0.02, f"Number of Nodes: {num_nodes}", ha="left", fontsize=12)
    plt.figtext(0.1, 0.01, f"Used Colors: {used_colors}", ha="left", fontsize=12)

    return ani
